"""
EdgeX exchange client implementation.
"""

import os
import asyncio
import json
import traceback
from decimal import Decimal
from typing import Dict, Any, List, Optional, Tuple
from edgex_sdk import Client, OrderSide, WebSocketManager, CancelOrderParams, GetOrderBookDepthParams, GetActiveOrderParams

from .base import BaseExchangeClient, OrderResult, OrderInfo
from helpers.logger import TradingLogger


class EdgeXClient(BaseExchangeClient):
    """EdgeX exchange client implementation."""

    def __init__(self, config: Dict[str, Any]):
        """Initialize EdgeX client."""
        super().__init__(config)

        # EdgeX credentials prefer per-bot config then env
        self.account_id = getattr(self.config, 'edgex_account_id', None) or os.getenv('EDGEX_ACCOUNT_ID')
        self.stark_private_key = getattr(self.config, 'edgex_stark_private_key', None) or os.getenv('EDGEX_STARK_PRIVATE_KEY')
        self.base_url = os.getenv('EDGEX_BASE_URL', 'https://pro.edgex.exchange')
        self.ws_url = os.getenv('EDGEX_WS_URL', 'wss://quote.edgex.exchange')

        if not self.account_id or not self.stark_private_key:
            raise ValueError("EDGEX_ACCOUNT_ID and EDGEX_STARK_PRIVATE_KEY are required (via per-bot config or env)")

        # Initialize EdgeX client using official SDK
        self.client = Client(
            base_url=self.base_url,
            account_id=int(self.account_id),
            stark_private_key=self.stark_private_key
        )

        # Initialize WebSocket manager using official SDK
        self.ws_manager = WebSocketManager(
            base_url=self.ws_url,
            account_id=int(self.account_id),
            stark_pri_key=self.stark_private_key
        )

        # Initialize logger using the same format as helpers
        self.logger = TradingLogger(exchange="edgex", ticker=self.config.ticker, log_to_console=False)
        try:
            # Populate logger context with account and (if derivable) address placeholder
            self.logger.set_account_id(str(self.account_id))
        except Exception:
            pass

        self._order_update_handler = None

    def _validate_config(self) -> None:
        """Validate EdgeX configuration."""
        # allow provided via config
        if not (getattr(self.config, 'edgex_account_id', None) and getattr(self.config, 'edgex_stark_private_key', None)):
            required_env_vars = ['EDGEX_ACCOUNT_ID', 'EDGEX_STARK_PRIVATE_KEY']
            missing_vars = [var for var in required_env_vars if not os.getenv(var)]
            if missing_vars:
                raise ValueError(f"Missing required credentials: {missing_vars}")

    async def connect(self) -> None:
        """Connect to EdgeX WebSocket."""
        self.ws_manager.connect_private()
        # Wait a moment for connection to establish
        await asyncio.sleep(2)

    async def disconnect(self) -> None:
        """Disconnect from EdgeX."""
        try:
            if hasattr(self, 'client') and self.client:
                await self.client.close()
            if hasattr(self, 'ws_manager'):
                self.ws_manager.disconnect_all()
        except Exception as e:
            self.logger.log(f"Error during EdgeX disconnect: {e}", "ERROR")

    def get_exchange_name(self) -> str:
        """Get the exchange name."""
        return "edgex"

    async def cancel_all_orders(self) -> None:
        """Cancel all active orders under the current account (EdgeX private API).

        Uses SDK: order.cancel_order(CancelOrderParams(contract_id=...)) which maps to
        POST /api/v1/private/order/cancelAllOrder with filterContractIdList.
        For 'all contracts', we enumerate active orders and cancel by each distinct contractId.
        """
        try:
            # 1) Get active orders to discover all contractIds with open orders
            params = GetActiveOrderParams(size="200", offset_data="")
            active = await self.client.get_active_orders(params)
            contract_ids = set()
            if active and isinstance(active, dict) and 'data' in active:
                for o in active['data'].get('dataList', []):
                    if isinstance(o, dict) and o.get('status') in ['OPEN', 'PARTIALLY_FILLED']:
                        cid = o.get('contractId')
                        if cid:
                            contract_ids.add(str(cid))

            if not contract_ids:
                self.logger.log("[EDGEX] no active orders to cancel", "INFO")
                return

            # 2) For each contract, call cancelAll
            count = 0
            self.logger.log(f"[EDGEX] will cancel orders for contracts: {sorted(list(contract_ids))}", "INFO")
            for cid in contract_ids:
                try:
                    await self.client.order.cancel_order(CancelOrderParams(contract_id=cid))
                    count += 1
                    self.logger.log(f"[EDGEX] cancelAll requested for contract {cid}", "INFO")
                except Exception as ce:
                    self.logger.log(f"[EDGEX] cancelAll failed for contract {cid}: {ce}", "ERROR")

            self.logger.log(f"[EDGEX] cancelAll finished for {count} contract(s)", "INFO")

        except Exception as e:
            self.logger.log(f"[EDGEX] cancel all failed: {e}", "ERROR")

    async def close_all_positions(self, side: Optional[str] = None) -> None:
        """Flatten all open positions using market orders via SDK create_market_order.

        Args:
            side: 'buy' or 'sell' to force close direction. If None, tries SELL then BUY.
        """
        try:
            # Optional: cancel all open orders first to avoid conflicts
            try:
                await self.cancel_all_orders()
            except Exception:
                pass

            self.logger.log("[EDGEX] fetching positions...", "INFO")
            positions_data = await self.client.get_account_positions()
            positions = []
            if positions_data and isinstance(positions_data, dict):
                positions = positions_data.get('data', {}).get('positionList', []) or []

            if not positions:
                self.logger.log("[EDGEX] no open positions to close", "INFO")
                return

            self.logger.log(f"[EDGEX] found {len(positions)} position entries", "INFO")
            closed_cnt = 0
            for p in positions:
                try:
                    size = Decimal(str(p.get('openSize', '0')))
                except Exception:
                    size = Decimal(0)
                if size <= 0:
                    self.logger.log(f"[EDGEX] skip position with zero size: {p}", "DEBUG")
                    continue
                contract_id = str(p.get('contractId'))
                if not contract_id:
                    self.logger.log(f"[EDGEX] skip: missing contractId in {p}", "ERROR")
                    continue

                # try derive direction from position if not forced
                derived_side = None
                try:
                    raw_side = (p.get('positionSide') or p.get('side') or p.get('direction') or '').upper()
                    if 'LONG' in raw_side or raw_side == 'BUY':
                        derived_side = 'SELL'  # to close long, sell
                    elif 'SHORT' in raw_side or raw_side == 'SELL':
                        derived_side = 'BUY'   # to close short, buy
                except Exception:
                    pass
                self.logger.log(f"[EDGEX] will close contract={contract_id} openSize={size} derivedSide={derived_side}", "INFO")

                async def try_market(side_str: str) -> bool:
                    try:
                        order_side = OrderSide.SELL if side_str.upper() == 'SELL' else OrderSide.BUY
                        res = await self.client.create_market_order(
                            contract_id=contract_id,
                            size=str(size),
                            side=order_side
                        )
                        self.logger.log(f"[EDGEX] market close {side_str.upper()} size={size} contract={contract_id} result={res}", "INFO")
                        return True
                    except Exception as e:
                        self.logger.log(f"[EDGEX] market close attempt {side_str} failed for contract {contract_id}: {e}", "ERROR")
                        return False

                if side in ('buy', 'sell'):
                    await try_market(side)
                else:
                    # Prefer derived side if available
                    if derived_side:
                        ok = await try_market(derived_side)
                        if not ok:
                            # fallback opposite
                            await try_market('BUY' if derived_side == 'SELL' else 'SELL')
                    else:
                        # Try SELL then BUY to cover unknown direction
                        ok = await try_market('SELL')
                        if not ok:
                            await try_market('BUY')
                closed_cnt += 1

            self.logger.log(f"[EDGEX] close-all finished for {closed_cnt} position(s)", "INFO")

        except Exception as e:
            self.logger.log(f"[EDGEX] close all positions failed: {e}", "ERROR")

    def setup_order_update_handler(self, handler) -> None:
        """Setup order update handler for WebSocket."""
        self._order_update_handler = handler

        def order_update_handler(message):
            """Handle order updates from WebSocket."""
            try:
                # Parse the message structure
                if isinstance(message, str):
                    message = json.loads(message)

                # Check if this is a trade-event with ORDER_UPDATE
                content = message.get("content", {})
                event = content.get("event", "")
                if event == "ORDER_UPDATE":
                    # Extract order data from the nested structure
                    data = content.get('data', {})
                    orders = data.get('order', [])

                    if orders and len(orders) > 0:
                        order = orders[0]  # Get the first order
                        order_id = order.get('id')
                        status = order.get('status')
                        side = order.get('side', '').lower()
                        filled_size = order.get('cumMatchSize')

                        if side == self.config.close_order_side:
                            order_type = "CLOSE"
                        else:
                            order_type = "OPEN"

                        # edgex returns TWO filled events for the same order; take the first one
                        if status == "FILLED" and len(data.get('collateral', [])):
                            return

                        # ignore canceled close orders
                        if status == "CANCELED" and order_type == "CLOSE":
                            return

                        # edgex returns partially filled events as "OPEN" orders
                        if status == "OPEN" and Decimal(filled_size) > 0:
                            status = "PARTIALLY_FILLED"

                        if status in ['OPEN', 'PARTIALLY_FILLED', 'FILLED', 'CANCELED']:
                            if self._order_update_handler:
                                self._order_update_handler({
                                    'order_id': order_id,
                                    'side': side,
                                    'order_type': order_type,
                                    'status': status,
                                    'size': order.get('size'),
                                    'price': order.get('price'),
                                    'contract_id': order.get('contractId'),
                                    'filled_size': filled_size
                                })

            except Exception as e:
                self.logger.log(f"Error handling order update: {e}", "ERROR")
                self.logger.log(f"Traceback: {traceback.format_exc()}", "ERROR")

        try:
            private_client = self.ws_manager.get_private_client()
            private_client.on_message("trade-event", order_update_handler)
        except Exception as e:
            self.logger.log(f"Could not add trade-event handler: {e}", "ERROR")

    async def place_open_order(self, contract_id: str, quantity: Decimal, direction: str) -> OrderResult:
        """Place an open order with EdgeX using official SDK with retry logic for POST_ONLY rejections."""
        max_retries = 15
        retry_count = 0

        while retry_count < max_retries:
            try:
                depth_params = GetOrderBookDepthParams(contract_id=contract_id, limit=15)
                order_book = await self.client.quote.get_order_book_depth(depth_params)

                # Handle the response format: {"code": "SUCCESS", "data": [{"asks": [...], "bids": [...]}]}
                if not isinstance(order_book, dict) or 'data' not in order_book:
                    return OrderResult(success=False, error_message='Unexpected order book response format')

                order_book_data = order_book['data']
                if not isinstance(order_book_data, list) or len(order_book_data) == 0:
                    return OrderResult(success=False, error_message='Order book data is not a valid list')

                # Get the first (and should be only) order book entry
                order_book_entry = order_book_data[0]
                if not isinstance(order_book_entry, dict):
                    return OrderResult(success=False, error_message='Order book entry is not a dict')

                # Extract bids and asks from the entry
                bids = order_book_entry.get('bids', [])
                asks = order_book_entry.get('asks', [])

                if not bids or not asks:
                    return OrderResult(success=False, error_message='No bid/ask data available')

                # Best bid is the highest price someone is willing to buy at
                best_bid = Decimal(bids[0]['price']) if bids and len(bids) > 0 else 0
                # Best ask is the lowest price someone is willing to sell at
                best_ask = Decimal(asks[0]['price']) if asks and len(asks) > 0 else 0

                if best_bid <= 0 or best_ask <= 0:
                    return OrderResult(success=False, error_message='Invalid bid/ask prices')

                if direction == 'buy':
                    # For buy orders, place slightly below best ask to ensure execution
                    order_price = best_ask - self.config.tick_size
                    side = OrderSide.BUY
                else:
                    # For sell orders, place slightly above best bid to ensure execution
                    order_price = best_bid + self.config.tick_size
                    side = OrderSide.SELL

                # Place the order using official SDK (post-only to ensure maker order)
                order_result = await self.client.create_limit_order(
                    contract_id=contract_id,
                    size=str(quantity),
                    price=str(self.round_to_tick(order_price)),
                    side=side,
                    post_only=True
                )

                if not order_result or 'data' not in order_result:
                    return OrderResult(success=False, error_message='Failed to place order')

                # Extract order ID from response
                order_id = order_result['data'].get('orderId')
                if not order_id:
                    return OrderResult(success=False, error_message='No order ID in response')

                # Check order status after a short delay to see if it was rejected
                await asyncio.sleep(0.01)
                order_info = await self.get_order_info(order_id)

                if order_info:
                    if order_info.status == 'CANCELED':
                        if retry_count < max_retries - 1:
                            retry_count += 1
                            continue
                        else:
                            return OrderResult(success=False, error_message=f'Order rejected after {max_retries} attempts')
                    elif order_info.status in ['OPEN', 'PARTIALLY_FILLED', 'FILLED']:
                        # Order successfully placed
                        return OrderResult(
                            success=True,
                            order_id=order_id,
                            side=side.value,
                            size=quantity,
                            price=order_price,
                            status=order_info.status
                        )
                    else:
                        return OrderResult(success=False, error_message=f'Unexpected order status: {order_info.status}')
                else:
                    # Assume order is successful if we can't get info
                    return OrderResult(
                        success=True,
                        order_id=order_id,
                        side=side.value,
                        size=quantity,
                        price=order_price
                    )

            except Exception as e:
                if retry_count < max_retries - 1:
                    retry_count += 1
                    await asyncio.sleep(0.1)  # Wait before retry
                    continue
                else:
                    return OrderResult(success=False, error_message=str(e))

        return OrderResult(success=False, error_message='Max retries exceeded')

    async def place_close_order(self, contract_id: str, quantity: Decimal, price: Decimal, side: str) -> OrderResult:
        """Place a close order with EdgeX using official SDK with retry logic for POST_ONLY rejections."""
        max_retries = 15
        retry_count = 0

        while retry_count < max_retries:
            try:
                # Get current market prices to adjust order price if needed
                depth_params = GetOrderBookDepthParams(contract_id=contract_id, limit=15)
                order_book = await self.client.quote.get_order_book_depth(depth_params)

                if not isinstance(order_book, dict) or 'data' not in order_book:
                    return OrderResult(success=False, error_message='Failed to get order book')

                order_book_data = order_book['data']
                if not isinstance(order_book_data, list) or len(order_book_data) == 0:
                    return OrderResult(success=False, error_message='Invalid order book data')

                # Get the first order book entry
                order_book_entry = order_book_data[0]
                bids = order_book_entry.get('bids', [])
                asks = order_book_entry.get('asks', [])

                if not bids or not asks:
                    return OrderResult(success=False, error_message='No bid/ask data available')

                # Get best bid and ask prices
                best_bid = Decimal(bids[0]['price']) if bids and len(bids) > 0 else 0
                best_ask = Decimal(asks[0]['price']) if asks and len(asks) > 0 else 0

                if best_bid <= 0 or best_ask <= 0:
                    return OrderResult(success=False, error_message='Invalid bid/ask prices')

                # Convert side string to OrderSide enum
                order_side = OrderSide.BUY if side.lower() == 'buy' else OrderSide.SELL

                # Adjust order price based on market conditions and side
                adjusted_price = price

                if side.lower() == 'sell':
                    # For sell orders, ensure price is above best bid to be a maker order
                    if price <= best_bid:
                        adjusted_price = best_bid + self.config.tick_size
                elif side.lower() == 'buy':
                    # For buy orders, ensure price is below best ask to be a maker order
                    if price >= best_ask:
                        adjusted_price = best_ask - self.config.tick_size

                adjusted_price = self.round_to_tick(adjusted_price)
                # Place the order using official SDK (post-only to avoid taker fees)
                order_result = await self.client.create_limit_order(
                    contract_id=contract_id,
                    size=str(quantity),
                    price=str(adjusted_price),
                    side=order_side,
                    post_only=True
                )

                if not order_result or 'data' not in order_result:
                    return OrderResult(success=False, error_message='Failed to place order')

                # Extract order ID from response
                order_id = order_result['data'].get('orderId')
                if not order_id:
                    return OrderResult(success=False, error_message='No order ID in response')

                # Check order status after a short delay to see if it was rejected
                await asyncio.sleep(0.01)
                order_info = await self.get_order_info(order_id)

                if order_info:
                    if order_info.status == 'CANCELED':
                        if retry_count < max_retries - 1:
                            retry_count += 1
                            continue
                        else:
                            return OrderResult(success=False, error_message=f'Close order rejected after {max_retries} attempts')
                    elif order_info.status in ['OPEN', 'PARTIALLY_FILLED', 'FILLED']:
                        # Order successfully placed
                        return OrderResult(
                            success=True,
                            order_id=order_id,
                            side=side,
                            size=quantity,
                            price=adjusted_price,
                            status=order_info.status
                        )
                    else:
                        return OrderResult(success=False, error_message=f'Unexpected close order status: {order_info.status}')
                else:
                    # Assume order is successful if we can't get info
                    return OrderResult(
                        success=True,
                        order_id=order_id,
                        side=side,
                        size=quantity,
                        price=adjusted_price
                    )

            except Exception as e:
                if retry_count < max_retries - 1:
                    retry_count += 1
                    await asyncio.sleep(0.1)  # Wait before retry
                    continue
                else:
                    return OrderResult(success=False, error_message=str(e))

        return OrderResult(success=False, error_message='Max retries exceeded for close order')

    async def cancel_order(self, order_id: str) -> OrderResult:
        """Cancel an order with EdgeX using official SDK."""
        try:
            # Create cancel parameters using official SDK
            cancel_params = CancelOrderParams(order_id=order_id)

            # Cancel the order using official SDK
            cancel_result = await self.client.cancel_order(cancel_params)

            if not cancel_result or 'data' not in cancel_result:
                return OrderResult(success=False, error_message='Failed to cancel order')

            return OrderResult(success=True)

        except Exception as e:
            return OrderResult(success=False, error_message=str(e))

    async def get_order_info(self, order_id: str) -> Optional[OrderInfo]:
        """Get order information from EdgeX using official SDK."""
        try:
            # Use the newly created get_order_by_id method
            order_result = await self.client.order.get_order_by_id(order_id_list=[order_id])

            if not order_result or 'data' not in order_result:
                return None

            # The API returns a list of orders, get the first (and should be only) one
            order_list = order_result['data']
            if order_list and len(order_list) > 0:
                order_data = order_list[0]
                return OrderInfo(
                    order_id=order_data.get('id', ''),
                    side=order_data.get('side', '').lower(),
                    size=Decimal(order_data.get('size', 0)),
                    price=Decimal(order_data.get('price', 0)),
                    status=order_data.get('status', ''),
                    filled_size=Decimal(order_data.get('cumMatchSize', 0)),
                    remaining_size=Decimal(order_data.get('size', 0)) - Decimal(order_data.get('cumMatchSize', 0))
                )

            return None

        except Exception:
            return None

    async def get_active_orders(self, contract_id: str) -> List[OrderInfo]:
        """Get active orders for a contract using official SDK."""
        try:
            # Get active orders using official SDK
            params = GetActiveOrderParams(size="100", offset_data="")
            active_orders = await self.client.get_active_orders(params)

            if not active_orders or 'data' not in active_orders:
                return []

            # Filter orders for the specific contract and ensure they are dictionaries
            # The API returns orders under 'dataList' key, not 'orderList'
            order_list = active_orders['data'].get('dataList', [])
            contract_orders = []

            for order in order_list:
                if isinstance(order, dict) and order.get('contractId') == contract_id:
                    contract_orders.append(OrderInfo(
                        order_id=order.get('id', ''),
                        side=order.get('side', '').lower(),
                        size=Decimal(order.get('size', 0)),
                        price=Decimal(order.get('price', 0)),
                        status=order.get('status', ''),
                        filled_size=Decimal(order.get('cumMatchSize', 0)),
                        remaining_size=Decimal(order.get('size', 0)) - Decimal(order.get('cumMatchSize', 0))
                    ))

            return contract_orders

        except Exception:
            return []

    async def get_account_positions(self) -> Decimal:
        """Get account positions using official SDK."""
        try:
            positions_data = await self.client.get_account_positions()
            if not positions_data or 'data' not in positions_data:
                self.logger.log("No positions or failed to get positions", "WARNING")
                position_amt = 0
            else:
                # The API returns positions under data.positionList
                positions = positions_data.get('data', {}).get('positionList', [])
                if positions:
                    # Find position for current contract
                    position = None
                    for p in positions:
                        if isinstance(p, dict) and p.get('contractId') == self.config.contract_id:
                            position = p
                            break

                    if position:
                        position_amt = abs(Decimal(position.get('openSize', 0)))
                    else:
                        position_amt = 0
                else:
                    position_amt = 0
            return position_amt
        except Exception:
            return 0

    async def get_contract_attributes(self) -> Tuple[str, Decimal]:
        """Get contract ID for a ticker."""
        ticker = self.config.ticker
        if len(ticker) == 0:
            self.logger.log("Ticker is empty", "ERROR")
            raise ValueError("Ticker is empty")

        response = await self.client.get_metadata()
        data = response.get('data', {})
        if not data:
            self.logger.log("Failed to get metadata", "ERROR")
            raise ValueError("Failed to get metadata")

        contract_list = data.get('contractList', [])
        if not contract_list:
            self.logger.log("Failed to get contract list", "ERROR")
            raise ValueError("Failed to get contract list")

        current_contract = None
        for c in contract_list:
            if c.get('contractName') == ticker+'USD':
                current_contract = c
                break

        if not current_contract:
            self.logger.log("Failed to get contract ID for ticker", "ERROR")
            raise ValueError("Failed to get contract ID for ticker")

        self.config.contract_id = current_contract.get('contractId')
        min_quantity = Decimal(current_contract.get('minOrderSize'))
        if self.config.quantity < min_quantity:
            self.logger.log(f"Order quantity is less than min quantity: {self.config.quantity} < {min_quantity}", "ERROR")
            raise ValueError(f"Order quantity is less than min quantity: {self.config.quantity} < {min_quantity}")

        self.config.tick_size = Decimal(current_contract.get('tickSize'))

        return self.config.contract_id, self.config.tick_size
